USER = 'embed'
EMBED_TBL = 'kg_sentence_bank_m3e'

if '__main__' == __name__:
    from python_nlp.kg.neo4j.config import NEO4J_CONFIG_CBLUE_DETAIL
    from neo4j import GraphDatabase
    from PyCmpltrtok.common import sep, md5
    import pymongo as pm
    from PyCmpltrtok.util_mongo import mongo_upsert, VALUE, mongo_get
    from python_nlp.embed.cblue_text2mongo import sentence2embed, VER_INT, VERSION_COL, EMBED_COL, USERNAME
    import torch
    from transformers import AutoModel, AutoTokenizer
    from transformers import AutoConfig

    # 连接Mongodb
    mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
    mdb = mongo['CBLUE']

    # 加载文本BERT
    dev = torch.device(0)
    if 0:
        # model_name = 'bert-base-chinese'  # 用模型名，需联网，需翻墙
        model_name = 'moka-ai/m3e-base'  # 用模型名，需联网，需翻墙
    else:
        # model_name = r'C:\Users\peter\.cache\huggingface\hub\models--bert-base-chinese\snapshots\8d2a91f91cc38c96bb8b4556ba70c392f8d5ee55'  # 使用下载后的路径，不需联网
        model_name = r'D:\_const\wsl\my_github\m3e-base'
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(dev)

    # 确定最大token长度
    xlen01 = tokenizer.model_max_length
    config = AutoConfig.from_pretrained(model_name)
    xlen02 = config.max_position_embeddings
    max_len = min(xlen01, xlen02)

    def _main():

        driver = GraphDatabase.driver(**NEO4J_CONFIG_CBLUE_DETAIL)
        with driver.session() as ss:
            for query in [
                "CREATE INDEX IF NOT EXISTS FOR (n:Node) ON (n.name)",
                "CREATE CONSTRAINT IF NOT EXISTS FOR (n:Node) REQUIRE (n.type, n.name) IS UNIQUE",
                "CREATE INDEX IF NOT EXISTS FOR (p:Predicate) ON (p.name)",
            ]:
                print(query)
                ss.run(query)

        # 确定类别
        if 0:
            with driver.session() as ss:
                query = "MATCH ()-[p:Predicate]-() RETURN DISTINCT p.name"
                print(query)
                xres = ss.run(query)
                xarr = []
                for i, xel in enumerate(xres):
                    xarr.append(xel[0])
                xarr = sorted(set(xarr))
        else:
            """
辅助治疗
临床表现
同义词
多发群体
发病性别倾向
病因
相关（导致）
病理分型
实验室检查
转移部位
辅助检查
多发季节
病史
药物治疗
高危因素
发病率
发病部位
阶段
病理生理
预后状况
手术治疗
发病机制
并发症
鉴别诊断
化疗
相关（症状）
影像学检查
预后生存率
内窥镜检查
预防
筛查
放射治疗
遗传因素
多发地区
发病年龄
外侵部位
相关（转化）
组织学检查
侵及周围组织转移的症状
治疗后症状
就诊科室
传播途径
死亡率
风险评估因素
        """
            xarr = [
                '同义词',
                '临床表现', '鉴别诊断', '相关（症状）'
                '辅助治疗', '药物治疗', '手术治疗', '放射治疗'
            ]

        with driver.session() as ss:
            query = "MATCH (n1:Node)-[p:Predicate]->(n2:Node) " \
                    "WHERE p.name IN %r " \
                    "RETURN DISTINCT n1.name, p.name, n2.name" % xarr
            print(query)
            xres = ss.run(query)
            for i, xel in enumerate(xres):
                # if i >= 1700:
                #     break
                if i % 50 == 0:
                    print(i)
                print('.', end='')

                # print(i, xel[0], xel[1], xel[2])
                xsents = [
                    xel[0] + '的' + xel[1] + '是' + xel[2] + '。',
                    xel[2] + '是' + xel[0] + '的' + xel[1] + '。',
                ]

                for xsent in xsents:
                    xmd5 = md5(xsent)

                    # 获取以有数据
                    xdict = mongo_get(mdb, EMBED_TBL, USERNAME, xmd5, only_value=False)

                    # 已经存在该句子的嵌入
                    if xdict and xdict.get(EMBED_COL, None):
                        # print(f'Skip {xmd5} {xtext}')
                        continue

                    # 做嵌入
                    xembed = sentence2embed(dev, model, tokenizer, xsent, max_seq_length=max_len)
                    print('+', end='')

                    # 入库
                    mongo_upsert(mdb, EMBED_TBL, USER, xmd5, {
                        VALUE: xsent,
                        EMBED_COL: xembed,
                        VERSION_COL: VER_INT,
                    })
            print()
            print('Over')

    _main()